import os
import pandas as pd


def process_and_resample_data(base_dir, output_dir):
    # 定义重新采样函数
    def resample_ohlc(data):
        ohlc_dict = {
            '开盘价': 'first',
            '最高价': 'max',
            '最低价': 'min',
            '收盘价': 'last',
            '成交量（手）': 'sum',
            '成交额（元）': 'sum'
        }
        return data.resample('15T', label='right', closed='right').apply(ohlc_dict)

    # 定义交易时间段
    def filter_trading_time(data):
        trading_times = (
                (data.index.time >= pd.to_datetime('09:31').time()) & (
                    data.index.time <= pd.to_datetime('11:30').time()) |
                (data.index.time >= pd.to_datetime('13:01').time()) & (
                            data.index.time <= pd.to_datetime('15:00').time())
        )
        return data[trading_times]

    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    # 遍历每年的文件
    for year in range(202401, 202406):
        # year_dir = os.path.join(base_dir, f'{year}年', f'{year}年')
        year_dir = os.path.join(base_dir, f'2024年', f'{year}')
        all_dfs = []

        for root, dirs, files in os.walk(year_dir):
            for file in files:
                if file.endswith('.csv'):
                    file_path = os.path.join(root, file)

                    # 读取分钟数据
                    df = pd.read_csv(file_path)

                    # 确保时间列被解析为datetime类型
                    df['日期'] = pd.to_datetime(df['日期'])

                    # 将时间列设为索引
                    df.set_index('日期', inplace=True)

                    # 提取代码列
                    stock_code = df['代码'].iloc[0]
                    converted_code = f"{stock_code.split('.')[1]}.{stock_code.split('.')[0]}"

                    # 重新采样数据
                    df_15min = resample_ohlc(df)

                    # 过滤掉非交易时间的数据
                    df_15min = filter_trading_time(df_15min)

                    # 保留原始数据的日期
                    df_15min = df_15min[df_15min.index.floor('T').isin(df.index.floor('T'))]

                    # 将时间戳重置为列
                    df_15min.reset_index(inplace=True)

                    # 添加代码列到时间戳后面
                    df_15min.insert(1, '代码', converted_code)

                    # 计算平均价格
                    df_15min['avg_price'] = round(df_15min['成交额（元）'] / df_15min['成交量（手）'] / 100, 2)

                    # 将处理后的数据添加到列表
                    all_dfs.append(df_15min)

        # 合并所有数据帧
        if all_dfs:
            combined_df = pd.concat(all_dfs)

            # 定义输出文件路径
            output_file_path = os.path.join(output_dir, f'15min_{year}.csv')

            # 保存转换后的数据
            combined_df.to_csv(output_file_path, index=False)

            print(f"Processed and saved: {output_file_path}")


# 调用函数处理数据
base_dir = 'F:\\Personal\\data\\行情数据\\'
output_dir = 'F:\\Personal\\data\\行情数据\\15min\\2024'
process_and_resample_data(base_dir, output_dir)
